import asyncio
import logging
import math
import os
import random
import shutil
import threading
import time
import urllib.parse
import uuid
from contextlib import contextmanager
from datetime import datetime, timedelta

import cv2
from sqlalchemy import func,create_engine

import task
from blls.base_bll import BaseBLL
from config import REDIS_DATABASE_KEY, MYSQL_USER, MYSQL_PASSWORD, MYSQL_HOST, MYSQL_PORT, MYSQL_DB
from libs.redis_wrapper import RedisWrapper
from models.annotation_models.annotation_model import Annotation
from models.data_models.data_model import Data
from models.label_models.label import Label
from models.project_models.project_model import Project
from models.setting_models.setting_model import TaskCategory
from models.user_models.system_user_model import SystemUsers
from utils.common_util import CommonUtil
from utils.guid_util import GUIDUtil
from dtos.result import Result
from enumobjects.general_enum import YesOrNoEnum
from task.util.file import expand_home
from pathlib import Path
import os.path as osp

from sqlalchemy.orm import scoped_session,sessionmaker
def my_same_file_diff_checker(*args, **kwargs):
    return False


shutil._samefile = my_same_file_diff_checker
class ProjectBLL(BaseBLL):
    # _session_factory = None
    # @classmethod
    # def _create_session_factory(cls):
    #     DATABASE_URI = f'mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DB}?charset=utf8mb4'
    #     engine = create_engine(DATABASE_URI)
    #     cls._session_factory = scoped_session(sessionmaker(bind=engine))
    #
    # @contextmanager
    # def _get_db_session(self):
    #     if self._session_factory is None:
    #         self._create_session_factory()
    #     session = self._session_factory()
    #     try:
    #         yield session
    #     finally:
    #         session.close()

    def export_dataset(self, project_id, export_format, export_dir, seg_mask_type, finalValues):
        MYSQL_PASSWORD_2 = urllib.parse.quote_plus(MYSQL_PASSWORD)
        DATABASE_URI = f'mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD_2}@{MYSQL_HOST}' \
                       f':{MYSQL_PORT}/{MYSQL_DB}?' \
                       f'charset=utf8mb4'
        engine = create_engine(DATABASE_URI)
        Session = scoped_session(sessionmaker(bind=engine))
        project_obj = self.get_one_obj_by_field('project_id', project_id, Project)
        if not project_obj:
            return Result.error_result(message="没有找到项目")
        task_category_obj = self.get_one_obj_by_field('task_category_id', project_obj.task_category_id,
                                                      TaskCategory)
        if not task_category_obj:
            return Result.error_result(message="没有找到项目类型")
        db_session = Session()
        handler = eval(task_category_obj.handler)(db_session, project_obj, is_export=True)
        if export_format is None or len(export_format) == 0:
            exporter = handler.default_exporter
        else:
            exporter = handler.exporters[export_format]
        export_dir = expand_home(export_dir)
        if not Path(export_dir).is_absolute():
            return Result.error_result(message="只支持绝对路径")

        if osp.exists(osp.join(export_dir, "paddlelabel.warning")):
            return Result.error_result(message="该文件夹被用作PaddleLabel的文件存储。请指定另一个文件夹进行导出")

        params_dict = {
            'export_dir': export_dir,
            'db_session':db_session
        }
        if seg_mask_type:
            params_dict['seg_mask_type'] = seg_mask_type
        if finalValues:
            params_dict['finalValues'] = finalValues

        # 异步调用 exporter
        import threading
        thread = threading.Thread(target=exporter, kwargs=params_dict)
        thread.start()
        return Result.success_result(message="导出任务已启动，将在后台处理")

    def split_dataset(self, project_id, split):
        project_obj = self.get_one_obj_by_field('project_id', project_id, Project)
        if not project_obj:
            return Result.error_result(message="未找到该项目")
        if list(split.keys()) != ['train', 'val', 'test']:
            logging.info("项目id：%s，数据分割类型：%s，数据分割失败" % (project_id, split))
            return Result.error_result(message="Request should provide train, validataion and test percentage")
        if sum(split.values()) != 100:
            return Result.error_result(message="三值总和应等于100")
        split_num = [0] * 3
        split_num[1] = split["train"] / 100
        split_num[2] = split["val"] / 100
        split = split_num
        for idx in range(1, 3):
            split[idx] += split[idx - 1]
        task_list = self.db_session.query(Data).filter(
            Data.is_delete == YesOrNoEnum.NO.value,
            Data.project_id == project_obj.project_id,
            Data.conform == 1,
        ).all()
        if not task_list:
            return Result.error_result(message="没有查到该项目组的任务")
        self.db_session.begin(subtransactions=True)
        try:
            split = [math.ceil(s * len(task_list)) for s in split]
            split.append(len(task_list))

            random.shuffle(task_list)
            for set_idx in range(3):
                for idx in range(split[set_idx], split[set_idx + 1]):
                    task_list[idx].set = set_idx
            self.db_session.commit()
            return Result.success_result(data={
                'train': split[1],
                'val': split[2] - split[1],
                'test': split[3] - split[2]
            })
        except Exception as ex:
            message = "划分数据失败，原因是：%s" % str(ex)
            logging.info(message)
            return Result.error_result(message=message)

    @staticmethod
    def import_dataset(self, project, data_dir=None, label_format=None):
        logging.info("进入import_dataset")
        data_dir = project.data_dir if data_dir is None else data_dir
        # logging.info(f'Importing dataset from {data_dir}')
        task_category = self.get_one_obj_by_field('task_category_id', project.task_category_id, TaskCategory)
        if not task_category:
            return Result.error_result(message="未找到该分类")
        handler = eval(task_category.handler)(session=self.db_session, project=project, data_dir=data_dir)
        # 2. choose importer. if specified, use importer for new_project.label_format, else use default_importer
        if label_format is None:
            label_format = None
            label_format = label_format if label_format else None
            # label_format = None if project.label_format is None else project.label_format
        if label_format not in [None, "default"] and label_format not in handler.importers.keys():
            return False, "label_format error"
        if label_format in [None, "default", ""]:
            importer = handler.default_importer
        else:
            importer = handler.importers[project.label_format]

        # 3. run import
        try:
            res, message = importer(data_dir)
            if not res:
                return res, message
        except Exception as ex:
            return False, "创建Data失败"
        return True, ""


    @staticmethod
    def makedirs(path, name):
        os.chdir(path)
        new_folder = name
        if not os.path.exists(new_folder):
            os.makedirs(new_folder)
        new_path = os.path.join(path, new_folder)
        return new_path

    def _process_images_multithreaded(self, input_path, output_path):
        img_list = [os.path.join(nm) for nm in os.listdir(input_path) if nm[-3:] in ['jpg', 'png', 'gif']]
        num_threads = min(len(img_list), 8)  # 最多使用8个线程
        if len(img_list) == 0:
            logging.info("No images found in the input directory.")
            return
        chunk_size = len(img_list) // num_threads
        threads = []

        def process_images(start, end):
            md5_list = set()  # 用于记录已处理的图片的MD5值
            for i in img_list[start:end]:
                md5 = CommonUtil.calculate_file_hash(os.path.join(input_path, i))
                if md5 in md5_list:
                    continue
                image = cv2.imread(os.path.join(input_path, i))
                cv2.imwrite(os.path.join(output_path, os.path.basename(i)), image)
                md5_list.add(md5)

        for i in range(num_threads):
            start = i * chunk_size
            end = (i + 1) * chunk_size if i < num_threads - 1 else len(img_list)
            thread = threading.Thread(target=process_images, args=(start, end))
            threads.append(thread)
            thread.start()

        for thread in threads:
            thread.join()

    @staticmethod
    def add_label(self, project_obj,name,color):
        project = project_obj
        if project is None:
            return Result.error_result(message="Project does not exist")
        new_label = Label(**{
            "project_id": project.project_id,
            "name": name,
            "color": color,
        })
        if new_label.id is None:
            ann = self.db_session.query(func.max(Label.id)).filter_by(project_id=new_label.project_id,
                                                                      is_delete=YesOrNoEnum.NO.value).one()[0]
            if ann is None:
                ann = 0
            new_label.id = ann + 1
        if new_label.color == "":
            colors = self.db_session.query(Label).with_entities(Label.color).filter(
                Label.project_id == new_label.project_id,
                Label.is_delete == YesOrNoEnum.NO.value
            ).all()
            colors = [c[0] for c in colors]
            new_label.color = CommonUtil.rand_hex_color(colors)

        # 3. cols must be unique within project
        cols = ["id", "name", "color"]
        labels = self.db_session.query(Label).filter(Label.project_id == new_label.project_id,
                                                     Label.is_delete == YesOrNoEnum.NO.value).all()
        rets, unique = CommonUtil.unique_within_project(labels, [new_label], cols)
        if not unique[0]:
            not_unique_cols = ", ".join([c for c, u in zip(cols, rets) if not u])
            return Result.error_result(message=f"Label {not_unique_cols} is not unique")
        label = Label(**{
            "guid": GUIDUtil.generate(),
            "id": new_label.id,
            "project_id": new_label.project_id,
            "name": new_label.name,
            "color": new_label.color,
            "create_time": CommonUtil.get_timestamp()
        })
        try:
            self.db_session.begin(subtransactions=True)
            self.db_session.add(label)
            self.db_session.commit()
        except Exception as ex:
            message = "添加错误，原因是:%s" % str(ex)
            logging.info(message)
            self.db_session.rollback()
            return Result.error_result(message=message)


    def create_project(self, **res_data):
        current_user_obj = self.get_current_user()
        if current_user_obj.role != 1:
            return Result.error_result(message='只有管理员可创建')
        if not res_data['name']:
            return Result.error_result(message='请填写项目名称')
        if self.db_session.query(Project).filter(Project.name == res_data['name'],
                                                  Project.is_delete == YesOrNoEnum.NO.value).first():
            return Result.error_result(message="已存在该项目名")
        if not res_data['data_dir']:
            return Result.error_result(message='请填写数据集路径')
        if self.db_session.query(Project).filter(Project.data_dir == res_data['data_dir'],
                                                  Project.is_delete == YesOrNoEnum.NO.value).first():
            return Result.error_result(message='已存在该路径')
        if not res_data['labels']:
            return Result.error_result(message='未添加标签')
        if not os.path.exists(res_data['data_dir']):
            return Result.error_result(message='当前目录不存在')
        if not os.path.isabs(res_data['data_dir']):
            return Result.error_result(message='数据集路径不是绝对路径')
        task_category_obj = self.db_session.query(TaskCategory).filter(
            TaskCategory.task_category_id == res_data['task_category_id'],
            TaskCategory.is_delete == YesOrNoEnum.NO.value).first()
        if not task_category_obj:
            return Result.error_result(message="未找到任务类别id")

        dir_path = res_data['data_dir']
        path = res_data['data_dir']
        name = "photo"
        new_path = self.makedirs(path, name)
        if not os.path.isabs(new_path):
            return Result.error_result('数据集路径不是绝对路径')

        try:
            _init_info = self._init_data({
                'name': res_data['name'],
                'description': res_data['description'],
                'task_category_id': res_data['task_category_id'],
                'data_dir': new_path,
            })
            project_info = Project(**_init_info)
            self.db_session.add(project_info)
            self.db_session.flush()
            project_obj = self.db_session.query(Project).filter_by(name=res_data['name']).first()
            if res_data['task_category_id'] != 1:
                for i in res_data['labels']:
                    self.add_label(self,project_obj=project_obj, name=i['label_name'], color=i['label_color'])
            else:
                for i in res_data['labels']:
                    self.add_label(self,project_obj=project_obj, name=i['label_name'], color="")
            # 并行处理图片
            self._process_images_multithreaded(dir_path, new_path)
            res_type, message = self.import_dataset(self, project=project_obj)
            if not res_type:
                return Result.error_result(message=message)
            return Result.success_result(message="创建项目成功")
        except Exception as ex:
            message = "创建项目失败，原因是：%s" % str(ex)
            logging.info(message)
            self.db_session.rollback()
            return Result.error_result(message=message)
    def state_project(self, **res_data):
        project_obj = self.get_one_obj_by_field('project_id', res_data['project_id'], Project)
        if not project_obj:
            return Result.error_result(message="未找到该项目信息")
        data_list = []
        ann_list = []
        for i in self.db_session.query(Data).filter_by(project_id=project_obj.project_id).all():
            data_list.append(i.data_id)
        for i in self.db_session.query(Annotation).filter(Annotation.is_delete == YesOrNoEnum.NO.value).all():
            ann_list.append(i.data_id)
        same_list = set(data_list).intersection(set(ann_list))
        state = 0
        if list(same_list) == data_list:
            state = 1
        if state == 1:
            try:
                self.db_session.begin()
                project_obj.state = state
                self.db_session.commit()
                return Result.success_result(message="项目状态更新成功")
            except Exception as ex:
                message = "项目状态更新失败，原因是：%s" % str(ex)
                logging.info(message)
                self.db_session.rollback()
                return Result.error_result(message=message)
        else:
            return Result.error_result(message="项目还未完成")

    def delete_project(self, **res_data):
        project_obj = self.db_session.query(Project).filter(
            Project.project_id == res_data['id'],
            Project.is_delete == YesOrNoEnum.NO.value
        ).first()
        if not project_obj:
            return Result.error_result(message="未找到要删除的信息")
        try:
            self.db_session.begin(subtransactions=True)
            project_obj.is_delete = YesOrNoEnum.YES.value
            data_obj = self.db_session.query(Data).filter(Data.project_id == project_obj.project_id).all()
            for i in data_obj:
                i.is_delete = YesOrNoEnum.YES.value
            label_obj = self.db_session.query(Label).filter(Label.project_id == project_obj.project_id).all()
            for i in label_obj:
                i.is_delete = YesOrNoEnum.YES.value
            annotation_obj = self.db_session.query(Annotation).filter(Annotation.project_id == project_obj.project_id).all()
            for i in annotation_obj:
                i.is_delete = YesOrNoEnum.YES.value
            self.db_session.commit()
            return Result.success_result(message="删除成功")
        except Exception as ex:
            message = "删除失败，原因是:%s" % str(ex)
            logging.info(message)
            self.db_session.rollback()
            return Result.error_result(message=message)

    def edit_project(self, **res_data):
        current_user_obj = self.get_current_user()
        project_obj = self.get_one_obj_by_field('project_id', res_data['id'], Project)
        if project_obj is None:
            return Result.error_result(message="未查到此项目")
        if project_obj.name == res_data['name']:
            return Result.error_result(message="名称已存在")
        self.db_session.begin(subtransactions=True)
        try:
            self.db_session.query(Project).filter(
                Project.project_id == res_data['id'],
                Project.is_delete == YesOrNoEnum.NO.value
            ).update({'name': res_data['name'],
                      "last_update_user_id": current_user_obj.id,
                      'description': res_data['description'],
                      "last_update_time": time.time(),
                      })
            if project_obj.task_category_id != 1:
                for i in res_data['labels']:
                    self.add_label(self,project_obj=project_obj, name=i['label_name'], color=i['label_color'])
            else:
                for i in res_data['labels']:
                    self.add_label(self,project_obj=project_obj, name=i['label_name'], color="")
            self.db_session.commit()
            return Result.success_result(message="修改项目成功")
        except Exception as ex:
            message = "修改项目失败，原因是：%s" % str(ex)
            logging.info(message)
            self.db_session.rollback()
            return Result.error_result(message=message)

    def detail_project(self, **res_data):
        """查询项目详情"""
        project_obj = self.db_session.query(Project).filter(Project.project_id == res_data['project_id'],Project.is_delete == YesOrNoEnum.NO.value).first()
        raw_data_list = [elem.decode('utf-8') for elem in
                   RedisWrapper.lrange(REDIS_DATABASE_KEY + str(res_data['project_id']) + str(res_data['username']), 0, -1)]
        data_list = list(set(raw_data_list))
        # data_list = data_list[::-1]
        # 取出最后 100条数据，如果不足100条则返回所有数据 用于indexdb中的数据初始化
        data_list = data_list[:100]
        if not project_obj:
            return Result.error_result(message="未找到该项目")
        data = {
            "id": project_obj.project_id,
            "guid": project_obj.guid,
            "description": project_obj.description,
            "name": project_obj.name,
            "task_category_id": project_obj.task_category_id,
            "data_dir": project_obj.data_dir,
            "label_format": None,
            "other_settings": None,
            "state": project_obj.state,
            "super_category_id": project_obj.super_category_id,
            "data_list": data_list
        }
        return Result.success_result(data=data)

    def list_project(self, **res_data):
        """查询各个任务所对应的项目,如目标检测任务对应的项目"""
        project_objs = (self.db_session.query(Project, TaskCategory).join(TaskCategory,Project.task_category_id == TaskCategory.task_category_id)
                        .filter(TaskCategory.is_delete == YesOrNoEnum.NO.value,
                                Project.is_delete == YesOrNoEnum.NO.value).order_by(Project.create_time.desc()))
        query, count = self.get_query_and_count(project_objs, res_data['page_size'], res_data['page_num'])
        project_list = query.all()
        data = {
            "page_num": res_data['page_num'],
            "page_size": res_data['page_size'],
            "count": count,
            "data_list": [{
                "created": CommonUtil.convert_timestamp_to_dt(i[0].create_time),
                "data_dir": i[0].data_dir,
                "description": i[0].description,
                "label_dir": None,
                "label_format": None,
                "labels": [],
                "modified": i[0].last_update_time,
                "name": i[0].name,
                "project_id": i[0].project_id,
                "state":"未完成" if i[0].state == 0 else "已完成",
                "task_category": {
                    "created": CommonUtil.convert_timestamp_to_dt(i[1].create_time),
                    "handler": i[1].handler,
                    "modified": i[1].last_update_time,
                    "name": i[1].name,
                    "task_category_id": i[1].task_category_id,
                },
                "super_category": i[0].super_category_id,
                "task_category_id": i[0].task_category_id,
                "upid": i[0].guid
            } for i in project_list]}
        return Result.success_result(data=data)

    def update_project(self, **res_data):
        project_obj = self.db_session.query(Project).filter(
            Project.project_id == res_data['project_id'],
            Project.is_delete == YesOrNoEnum.NO.value).first()
        if not project_obj:
            return Result.error_result(message="未查到此项目")
        if not os.path.isabs(res_data['path']):
            return Result.error_result(message='数据集路径不是绝对路径')
        if not os.path.exists(res_data['path']):
            return Result.error_result(message='当前目录不存在')

        data_objs = self.db_session.query(Data).filter(Data.project_id == res_data['project_id'],
                                                       Data.is_delete == YesOrNoEnum.NO.value).all()
        md5_lists = [i.md5 for i in data_objs]

        input_path = res_data['path']
        output_path = project_obj.data_dir

        self.db_session.begin()
        try:
            project_obj.state = 0
            self.db_session.commit()  # 提交事务以确保状态更新生效
        except Exception as ex:
            self.db_session.rollback()
            return Result.error_result(message=f"更新项目状态失败: {str(ex)}")

        img_list = [os.path.join(input_path, nm) for nm in os.listdir(input_path) if nm.endswith(('jpg', 'png', 'gif'))]
        response = Result.success_result(message="正在更新")

        thread = threading.Thread(target=ProjectBLL._async_update_project,
                                  args=(self, img_list, input_path, output_path, res_data['project_id'], md5_lists))
        thread.start()

        return response
        # time_1 = time.time()
        # for i in img_list:
        #     md5 = CommonUtil.calculate_file_hash(os.path.join(input_path, i))
        #     if md5 in md5_lists:
        #         continue
        #     else:
        #         image = cv2.imread(os.path.join(input_path, i))
        #         cv2.imwrite(os.path.join(output_path, os.path.basename(i)), image)
        # logging.info("cv2.imwrite耗时：%f",time.time()-time_1)
        # try:
        #     time_2 = time.time()
        #     res_type, message = self.import_dataset(self, project=project_obj)
        #     logging.info("import_time耗时：%f",time.time()-time_2)
        #     if not res_type:
        #         return Result.error_result(message=message)
        #     return Result.success_result(message="项目更新成功")
        #
        # except Exception as ex:
        #     message = "项目更新失败，原因是：%s" % str(ex)
        #     logging.info(message)
        #     self.db_session.rollback()
        #     return Result.error_result(message=message)

    @staticmethod
    def _async_update_project(self, img_list, input_path, output_path, project_id, md5_lists):
        logging.info("进入_async_update_project")

        # 重新查询 project_obj 以确保其绑定到当前会话
        project_obj = self.db_session.query(Project).filter(
            Project.project_id == project_id,
            Project.is_delete == YesOrNoEnum.NO.value).first()
        if not project_obj:
            logging.error("未查到此项目")
            return

        for i in img_list:
            md5 = CommonUtil.calculate_file_hash(os.path.join(input_path, i))
            if md5 in md5_lists:
                continue
            else:
                image = cv2.imread(os.path.join(input_path, i))
                cv2.imwrite(os.path.join(output_path, os.path.basename(i)), image)
                logging.info("写入图片：%s", os.path.basename(i))

        try:
            res_type, message = ProjectBLL.import_dataset(self, project=project_obj)
            if not res_type:
                logging.error(message)
                self.db_session.rollback()
            else:
                logging.info("项目更新成功")
                # 向 Redis 插入一个值表示更新完成
                user_id = self.get_current_user_id()
                key = f"project_update:{project_id}:{user_id}"
                RedisWrapper.set(key, "completed", ex=21600)
        except Exception as ex:
            message = "项目更新失败,原因是：%s" % str(ex)
            logging.error(message)
            self.db_session.rollback()
    def version(self, **res_data):
        data = "1.0.1"
        return Result.success_result(data=data)

    def ann_quantity(self, **res_data):
        # 将字符串转换为 datetime 对象
        start_time = datetime.strptime(res_data['start_time'], "%Y-%m-%d %H:%M")
        end_time = datetime.strptime(res_data['end_time'], "%Y-%m-%d %H:%M")
        # 将 datetime 对象转换为时间戳
        start_timestamp = time.mktime(start_time.timetuple())
        end_timestamp = time.mktime(end_time.timetuple())
        username = res_data['username']
        if not res_data['username']:
            result = self.db_session.query(
                SystemUsers.username,
                func.count(func.distinct(Annotation.data_id)).label('annotated_data_count')
            ).join(
                Annotation, SystemUsers.id == Annotation.create_user_id
            ).filter(
                Annotation.project_id == res_data['project_id'],
                Annotation.is_delete == YesOrNoEnum.NO.value,
                Annotation.create_time.between(int(start_timestamp), int(end_timestamp))
            ).group_by(
                SystemUsers.username
            ).all()
        else:
            user_id = self.db_session.query(SystemUsers).filter(SystemUsers.username == username).first().id
            result = self.db_session.query(
                SystemUsers.username,
                func.count(func.distinct(Annotation.data_id)).label('annotated_data_count')
            ).join(
                Annotation, SystemUsers.id == Annotation.create_user_id
            ).filter(
                Annotation.project_id == res_data['project_id'],
                Annotation.is_delete == YesOrNoEnum.NO.value,
                Annotation.create_user_id == user_id,
                Annotation.create_time.between(int(start_timestamp), int(end_timestamp))
            ).group_by(
                SystemUsers.username
            ).all()
        try:
            result_list = [{"username": row.username, "annotated_data_count": row.annotated_data_count} for row in result]
            return Result.success_result(data=result_list, message="查询成功")
        except Exception as ex:
            message = f"查询失败，原因是：{str(ex)}"
            logging.error(message)
            self.db_session.rollback()
            return Result.error_result(message=message)


    def conform(self, **res_data):
        data_obj = self.db_session.query(Data).filter(Data.data_id == res_data['data_id'],Data.is_delete == YesOrNoEnum.NO.value).first()
        if not data_obj:
            return Result.error_result(message="设置失败,未查到该data")
        self.db_session.begin()
        if res_data['status'] == 1:
            data_obj.conform = 1
        else:
            data_obj.conform = 0
        self.db_session.commit()
        return Result.success_result(data=data_obj.conform)
